[flang][cuda] Set correct offsets for multiple variables in dynamic shared memory #131674
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
llvmbot
added
flang
Member
|
@llvm/pr-subscribers-flang-fir-hlfir Author: Valentin Clement (バレンタイン クレメン) (clementval) ChangesCompute the dynamic size of variables to set the offsets correctly for multiple variables pointing to the dynamic shared memory. The first variables gets offset 0. Full diff: https://github.com/llvm/llvm-project/pull/131674.diff 2 Files Affected:
diff --git a/flang/lib/Optimizer/Transforms/CUFComputeSharedMemoryOffsetsAndSize.cpp b/flang/lib/Optimizer/Transforms/CUFComputeSharedMemoryOffsetsAndSize.cpp
index aec3ea294ac6c..dcb5f42902ee6 100644
--- a/flang/lib/Optimizer/Transforms/CUFComputeSharedMemoryOffsetsAndSize.cpp
+++ b/flang/lib/Optimizer/Transforms/CUFComputeSharedMemoryOffsetsAndSize.cpp
@@ -58,11 +58,13 @@ struct CUFComputeSharedMemoryOffsetsAndSize
auto gpuMod = cuf::getOrCreateGPUModule(mod, symTab);
mlir::Type i8Ty = builder.getI8Type();
mlir::Type i32Ty = builder.getI32Type();
+ mlir::Type idxTy = builder.getIndexType();
for (auto funcOp : gpuMod.getOps<mlir::gpu::GPUFuncOp>()) {
unsigned nbDynamicSharedVariables = 0;
unsigned nbStaticSharedVariables = 0;
uint64_t sharedMemSize = 0;
unsigned short alignment = 0;
+ mlir::Value crtDynOffset;
// Go over each shared memory operation and compute their start offset and
// the size and alignment of the global to be generated if all variables
@@ -73,16 +75,30 @@ struct CUFComputeSharedMemoryOffsetsAndSize
builder.setInsertionPoint(sharedOp);
if (fir::hasDynamicSize(sharedOp.getInType())) {
mlir::Type ty = sharedOp.getInType();
- // getTypeSizeAndAlignmentOrCrash will crash trying to compute the
- // size of an array with dynamic size. Just get the alignment to
- // create the global.
if (auto seqTy = mlir::dyn_cast<fir::SequenceType>(ty))
ty = seqTy.getEleTy();
unsigned short align = dl->getTypeABIAlignment(ty);
- ++nbDynamicSharedVariables;
- mlir::Value zero = builder.createIntegerConstant(loc, i32Ty, 0);
- sharedOp.getOffsetMutable().assign(zero);
alignment = std::max(alignment, align);
+ uint64_t tySize = dl->getTypeSize(ty);
+ ++nbDynamicSharedVariables;
+ if (crtDynOffset) {
+ sharedOp.getOffsetMutable().assign(
+ builder.createConvert(loc, i32Ty, crtDynOffset));
+ } else {
+ mlir::Value zero = builder.createIntegerConstant(loc, i32Ty, 0);
+ sharedOp.getOffsetMutable().assign(zero);
+ }
+
+ mlir::Value dynSize =
+ builder.createIntegerConstant(loc, idxTy, tySize);
+ for (auto extent : sharedOp.getShape())
+ dynSize = builder.create<mlir::arith::MulIOp>(loc, dynSize, extent);
+ if (crtDynOffset)
+ crtDynOffset =
+ builder.create<mlir::arith::AddIOp>(loc, crtDynOffset, dynSize);
+ else
+ crtDynOffset = dynSize;
+
continue;
}
auto [size, align] = fir::getTypeSizeAndAlignmentOrCrash(
diff --git a/flang/test/Fir/CUDA/cuda-shared-offset.mlir b/flang/test/Fir/CUDA/cuda-shared-offset.mlir
index 1eea75c802204..5e9aac4e71438 100644
--- a/flang/test/Fir/CUDA/cuda-shared-offset.mlir
+++ b/flang/test/Fir/CUDA/cuda-shared-offset.mlir
@@ -54,3 +54,56 @@ module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<!llvm.ptr, dense<
// CHECK: fir.global internal @_QPshared_static__shared_mem(dense<0> : vector<28xi8>) {alignment = 8 : i64, data_attr = #cuf.cuda<shared>} : !fir.array<28xi8>
// CHECK: }
// CHECK: }
+
+// -----
+
+module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<!llvm.ptr, dense<64> : vector<4xi64>>, #dlti.dl_entry<!llvm.ptr<271>, dense<32> : vector<4xi64>>, #dlti.dl_entry<!llvm.ptr<270>, dense<32> : vector<4xi64>>, #dlti.dl_entry<f128, dense<128> : vector<2xi64>>, #dlti.dl_entry<f64, dense<64> : vector<2xi64>>, #dlti.dl_entry<f80, dense<128> : vector<2xi64>>, #dlti.dl_entry<f16, dense<16> : vector<2xi64>>, #dlti.dl_entry<i32, dense<32> : vector<2xi64>>, #dlti.dl_entry<i16, dense<16> : vector<2xi64>>, #dlti.dl_entry<i128, dense<128> : vector<2xi64>>, #dlti.dl_entry<i8, dense<8> : vector<2xi64>>, #dlti.dl_entry<!llvm.ptr<272>, dense<64> : vector<4xi64>>, #dlti.dl_entry<i64, dense<64> : vector<2xi64>>, #dlti.dl_entry<i1, dense<8> : vector<2xi64>>, #dlti.dl_entry<"dlti.endianness", "little">, #dlti.dl_entry<"dlti.stack_alignment", 128 : i64>>, fir.defaultkind = "a1c4d8i4l4r4", fir.kindmap = "", gpu.container_module, llvm.data_layout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32:64-S128", llvm.ident = "flang version 20.0.0 (https://github.com/llvm/llvm-project.git cae351f3453a0a26ec8eb2ddaf773c24a29d929e)", llvm.target_triple = "x86_64-unknown-linux-gnu"} {
+ gpu.module @cuda_device_mod {
+ gpu.func @_QMmPshareddyn(%arg0: !fir.box<!fir.array<?x?xi32>> {cuf.data_attr = #cuf.cuda<device>, fir.bindc_name = "a"}, %arg1: !fir.box<!fir.array<?x?xi32>> {cuf.data_attr = #cuf.cuda<device>, fir.bindc_name = "b"}, %arg2: i32 {fir.bindc_name = "k"}) attributes {cuf.proc_attr = #cuf.cuda_proc<global>} {
+ %c1_i32 = arith.constant 1 : i32
+ %c2_i32 = arith.constant 2 : i32
+ %c0 = arith.constant 0 : index
+ %5 = fir.address_of(@_QM__fortran_builtinsE__builtin_blockdim) : !fir.ref<!fir.type<_QM__fortran_builtinsT__builtin_dim3{x:i32,y:i32,z:i32}>>
+ %6 = fir.declare %5 {uniq_name = "_QM__fortran_builtinsE__builtin_blockdim"} : (!fir.ref<!fir.type<_QM__fortran_builtinsT__builtin_dim3{x:i32,y:i32,z:i32}>>) -> !fir.ref<!fir.type<_QM__fortran_builtinsT__builtin_dim3{x:i32,y:i32,z:i32}>>
+ %15 = fir.alloca i32
+ %16 = fir.declare %15 {fortran_attrs = #fir.var_attrs<value>, uniq_name = "_QMmFss1Ek"} : (!fir.ref<i32>) -> !fir.ref<i32>
+ %27 = fir.coordinate_of %6, x : (!fir.ref<!fir.type<_QM__fortran_builtinsT__builtin_dim3{x:i32,y:i32,z:i32}>>) -> !fir.ref<i32>
+ %28 = fir.load %27 : !fir.ref<i32>
+ %29 = fir.convert %28 : (i32) -> i64
+ %30 = fir.convert %29 : (i64) -> index
+ %31 = arith.cmpi sgt, %30, %c0 : index
+ %32 = arith.select %31, %30, %c0 : index
+ %33 = fir.coordinate_of %6, y : (!fir.ref<!fir.type<_QM__fortran_builtinsT__builtin_dim3{x:i32,y:i32,z:i32}>>) -> !fir.ref<i32>
+ %34 = fir.load %33 : !fir.ref<i32>
+ %35 = fir.convert %34 : (i32) -> i64
+ %36 = fir.convert %35 : (i64) -> index
+ %37 = arith.cmpi sgt, %36, %c0 : index
+ %38 = arith.select %37, %36, %c0 : index
+ %39 = cuf.shared_memory !fir.array<?x?xi32>, %32, %38 : index, index {bindc_name = "s1", uniq_name = "_QMmFss1Es1"} -> !fir.ref<!fir.array<?x?xi32>>
+ %40 = fir.shape %32, %38 : (index, index) -> !fir.shape<2>
+ %41 = fir.declare %39(%40) {data_attr = #cuf.cuda<shared>, uniq_name = "_QMmFss1Es1"} : (!fir.ref<!fir.array<?x?xi32>>, !fir.shape<2>) -> !fir.ref<!fir.array<?x?xi32>>
+ %42 = fir.load %16 : !fir.ref<i32>
+ %43 = arith.muli %42, %c2_i32 : i32
+ %44 = fir.convert %43 : (i32) -> i64
+ %45 = fir.convert %44 : (i64) -> index
+ %46 = arith.cmpi sgt, %45, %c0 : index
+ %47 = arith.select %46, %45, %c0 : index
+ %48 = fir.load %16 : !fir.ref<i32>
+ %49 = fir.convert %48 : (i32) -> i64
+ %50 = fir.convert %49 : (i64) -> index
+ %51 = arith.cmpi sgt, %50, %c0 : index
+ %52 = arith.select %51, %50, %c0 : index
+ %53 = cuf.shared_memory !fir.array<?x?xi32>, %47, %52 : index, index {bindc_name = "s2", uniq_name = "_QMmFss1Es2"} -> !fir.ref<!fir.array<?x?xi32>>
+ gpu.return
+ }
+ }
+}
+
+// CHECK: gpu.func @_QMmPshareddyn(%arg0: !fir.box<!fir.array<?x?xi32>> {cuf.data_attr = #cuf.cuda<device>, fir.bindc_name = "a"}, %arg1: !fir.box<!fir.array<?x?xi32>> {cuf.data_attr = #cuf.cuda<device>, fir.bindc_name = "b"}, %arg2: i32 {fir.bindc_name = "k"}) attributes {cuf.proc_attr = #cuf.cuda_proc<global>} {
+// CHECK: %[[EXTENT0:.*]] = arith.select
+// CHECK: %[[EXTENT1:.*]] = arith.select
+// CHECK: %[[SIZE_EXTENT:.*]] = arith.muli %c4{{.*}}, %[[EXTENT0]] : index
+// CHECK: %[[DYNSIZE:.*]] = arith.muli %[[SIZE_EXTENT]], %[[EXTENT1]] : index
+// CHECK: cuf.shared_memory[%c0{{.*}} : i32] !fir.array<?x?xi32>, %9, %15 : index, index {bindc_name = "s1", uniq_name = "_QMmFss1Es1"} -> !fir.ref<!fir.array<?x?xi32>>
+// CHECK: %[[CONV_DYNSIZE:.*]] = fir.convert %[[DYNSIZE]] : (index) -> i32
+// CHECK: cuf.shared_memory[%[[CONV_DYNSIZE]] : i32] !fir.array<?x?xi32>, %26, %31 : index, index {bindc_name = "s2", uniq_name = "_QMmFss1Es2"} -> !fir.ref<!fir.array<?x?xi32>>
|
wangzpgi
approved these changes
Mar 17, 2025
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Compute the dynamic size of variables to set the offsets correctly for multiple variables pointing to the dynamic shared memory. The first variables gets offset 0.